from customize.get_ollama import GetOllama
from langchain import hub
from langchain_community.agent_toolkits import load_tools
from langchain_my_tools import bocha_tool

from langchain.agents import AgentExecutor
from langchain.agents.output_parsers import ReActSingleInputOutputParser
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.tools.render import render_text_description

llm = GetOllama(ip=GetOllama.ailab_linux_ip, model_name="qwen2.5:14b", model_type=1, temperature=0)()
llm_with_stop = llm.bind(stop=['\nObservation'])
prompt = hub.pull("hwchase17/react")
tools = load_tools.load_tools(["llm-math"], llm=llm)
#tools = load_tools.load_tools(["ddg-search", "llm-math"], llm=llm)
# tools.append(bocha_tool)
tools.append(bocha_tool)
prompt = prompt.partial(
    tools=render_text_description(tools),
    tool_names=", ".join([t.name for t in tools])
)
print(prompt)
agent = (
    {
        "input": lambda x: x['input'],
        "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),
    }
    | prompt
    | llm_with_stop
    | ReActSingleInputOutputParser()
)


agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
agent_executor.invoke({"input": "当前广东生猪价格是多少？如果我有生猪300头，每头平均重100千克，现在卖出可以有多少收入？"})
